import os
import cv2
import numpy as np

def create_large_image_and_yolo_labels(input_folder, output_image_folder, output_label_folder, class_id):
    image_files = [f for f in os.listdir(input_folder) if f.endswith(('.jpg', '.png', '.jpeg'))]
    image_files.sort()
    image_count = 0
    large_image_count = 0

    def get_class_name(class_id):
        if class_id == 0:
            return "GranularSludge"
        elif class_id == 1:
            return "Epistylis"
        elif class_id == 2:
            return "Floc"
        elif class_id == 3:
            return "Rotifer"
        elif class_id == 4:
            return "FilamentousBacteria"

    class_name = get_class_name(class_id)

    while image_count < len(image_files):
        large_image = np.zeros((1024, 1280, 3), dtype=np.uint8)
        yolo_labels = []

        for i in range(8):
            for j in range(10):
                if image_count >= len(image_files):
                    break

                img_path = os.path.join(input_folder, image_files[image_count])
                img = cv2.imread(img_path)
                img = cv2.resize(img, (128, 128))
                large_image[i*128:(i+1)*128, j*128:(j+1)*128] = img

                x_center = (j * 128 + 64) / 1280
                y_center = (i * 128 + 64) / 1024
                yolo_labels.append(f"{class_id} {x_center} {y_center} 0.1 0.125")

                image_count += 1

        cv2.imwrite(os.path.join(output_image_folder, f"large_image_{class_name}_{large_image_count}.bmp"), large_image)
        with open(os.path.join(output_label_folder, f"large_image_{class_name}_{large_image_count}.txt"), 'w') as f:
            f.write("\n".join(yolo_labels))

        large_image_count += 1

# input_folder = "D:/KLG/FINAL/code/stylegan3-main/stylegan3-main/out"
# output_image_folder = "D:/KLG/FINAL/code/micro_split/ganFB/images"
# output_label_folder = "D:/KLG/FINAL/code/micro_split/ganFB/labels"
# create_large_image_and_yolo_labels(input_folder, output_image_folder, output_label_folder)
